#include <memory>
#include <rclcpp/rclcpp.hpp>
#include <moveit/move_group_interface/move_group_interface.h>
#include <moveit/planning_scene_interface/planning_scene_interface.h>
#include "tf2/LinearMath/Quaternion.h"
#include <thread>
#include <chrono>
#include <tf2_ros/buffer.h>
#include <tf2_ros/transform_listener.h>

using namespace std::chrono_literals;

class ArucoPickPlaceDemo : public rclcpp::Node {
public:
  ArucoPickPlaceDemo() : Node("aruco_pick_place_demo",
      rclcpp::NodeOptions().automatically_declare_parameters_from_overrides(true)),
      tf_buffer_(std::make_unique<tf2_ros::Buffer>(this->get_clock())),
      tf_listener_(std::make_shared<tf2_ros::TransformListener>(*tf_buffer_)),
      arm_(std::make_shared<moveit::planning_interface::MoveGroupInterface>(
          shared_from_this(), "ur_manipulator")),
      gripper_(std::make_shared<moveit::planning_interface::MoveGroupInterface>(
          shared_from_this(), "gripper"))
  {
    // 初始化运动参数
    arm_->setMaxVelocityScalingFactor(0.5);
    arm_->setGoalPositionTolerance(0.01);
    arm_->setGoalOrientationTolerance(0.05);
  }

  bool run() {
    try {
      // 1. 回到home位置
      if (!moveToHome()) return false;

      // 2. 获取aruco marker位姿
      geometry_msgs::msg::TransformStamped transform;
      if (!getArucoTransform(transform)) return false;

      // 3. 执行抓取流程
      return executePickPlace(transform);

    } catch (const std::exception &e) {
      RCLCPP_ERROR(this->get_logger(), "Exception caught: %s", e.what());
      return false;
    }
  }

private:
  bool moveToHome() {
    RCLCPP_INFO(this->get_logger(), "Moving to home position");
    arm_->setNamedTarget("home");
    auto const result = arm_->move();
    return result == moveit::core::MoveItErrorCode::SUCCESS;
  }

  bool getArucoTransform(geometry_msgs::msg::TransformStamped &transform) {
    try {
      rclcpp::Time now = this->get_clock()->now();
      transform = tf_buffer_->lookupTransform(
        arm_->getPlanningFrame(), "aruco_marker_frame",now
      );
      return true;
    } catch (tf2::TransformException &ex) {
      RCLCPP_ERROR(this->get_logger(), "TF lookup failed: %s", ex.what());
      return false;
    }
  }

  bool executePickPlace(const geometry_msgs::msg::TransformStamped &transform) {
    // 获取home点的当前姿态
    geometry_msgs::msg::PoseStamped current_pose = arm_->getCurrentPose();
    
    // 提取marker的姿态和位置
    auto marker_pose = transform.transform;
    tf2::Quaternion marker_quat(
      marker_pose.rotation.x,
      marker_pose.rotation.y,
      marker_pose.rotation.z,
      marker_pose.rotation.w);

    // 计算末端执行器姿态 (保持当前姿态，仅更新z轴旋转)
    geometry_msgs::msg::PoseStamped target_pose;
    target_pose.header.frame_id = arm_->getPlanningFrame();
    target_pose.pose.position.x = marker_pose.translation.x;
    target_pose.pose.position.y = marker_pose.translation.y;
    target_pose.pose.position.z = marker_pose.translation.z;
    
    // 获取marker的z轴旋转角度
    double roll, pitch, yaw;
    tf2::Matrix3x3(marker_quat).getRPY(roll, pitch, yaw);
    
    // 获取当前姿态的roll和pitch
    tf2::Quaternion current_quat(
      current_pose.pose.orientation.x,
      current_pose.pose.orientation.y,
      current_pose.pose.orientation.z,
      current_pose.pose.orientation.w);
    
    double current_roll, current_pitch, current_yaw;
    tf2::Matrix3x3(current_quat).getRPY(current_roll, current_pitch, current_yaw);
    
    // 保持当前roll和pitch，仅使用marker的yaw
    tf2::Quaternion end_effector_quat;
    end_effector_quat.setRPY(current_roll, current_pitch, yaw);
    
    target_pose.pose.orientation.x = end_effector_quat.x();
    target_pose.pose.orientation.y = end_effector_quat.y();
    target_pose.pose.orientation.z = end_effector_quat.z();
    target_pose.pose.orientation.w = end_effector_quat.w();

    // 2. 移动到物体上方
    RCLCPP_INFO(this->get_logger(), "Moving above the object");
    target_pose.pose.position.z += 0.5;
    arm_->setPoseTarget(target_pose);
    if (!arm_->move()) return false;

    // 3. 移动到物体处
    RCLCPP_INFO(this->get_logger(), "Moving to the object");
    target_pose.pose.position.z -= 0.2;
    arm_->setPoseTarget(target_pose);
    if (!arm_->move()) return false;

    // 4. 闭合夹爪
    RCLCPP_INFO(this->get_logger(), "Closing gripper");
    gripper_->setNamedTarget("close");
    if (!gripper_->move()) return false;

    // 5. 回到物体上方
    RCLCPP_INFO(this->get_logger(), "Moving above the object");
    target_pose.pose.position.z += 0.2;
    arm_->setPoseTarget(target_pose);
    if (!arm_->move()) return false;

    // 6. 返回home位置
    return moveToHome();
  }

  std::unique_ptr<tf2_ros::Buffer> tf_buffer_;
  std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
  std::shared_ptr<moveit::planning_interface::MoveGroupInterface> arm_;
  std::shared_ptr<moveit::planning_interface::MoveGroupInterface> gripper_;
};

int main(int argc, char * argv[])
{
  rclcpp::init(argc, argv);
  
  // 创建单线程执行器
  rclcpp::executors::SingleThreadedExecutor executor;
  auto node = std::make_shared<ArucoPickPlaceDemo>();
  executor.add_node(node);
  auto spinner = std::thread([&executor]() { executor.spin(); });

  // 运行演示程序
  bool success = node->run();

  rclcpp::shutdown();
  spinner.join();
  return success ? 0 : 1;
}